# utils/init_utils.py
from config import (
    local_model_dir,
    CHROMA_PERSIST_DIR,
    MYSQL_USER,
    MYSQL_PASSWORD,
    MYSQL_HOST,
    MYSQL_PORT,
    MYSQL_DB
)
from transformers import AutoTokenizer, AutoModel
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
import chromadb

def get_tokenizer_and_model():
    """初始化文本向量化模型"""
    try:
        tokenizer = AutoTokenizer.from_pretrained(local_model_dir)
        model = AutoModel.from_pretrained(local_model_dir)
        print(f"成功加载本地模型：{local_model_dir}")
    except Exception as e:
        # 回退到远程模型
        tokenizer = AutoTokenizer.from_pretrained("uer/chinese_roberta_L-2_H-128")
        model = AutoModel.from_pretrained("uer/chinese_roberta_L-2_H-128")
        print("加载远程模型：uer/chinese_roberta_L-2_H-128")
    return tokenizer, model

def get_chroma_client():
    """初始化 ChromaDB 客户端"""
    return chromadb.Client(
        settings=chromadb.config.Settings(
            persist_directory=CHROMA_PERSIST_DIR,
            anonymized_telemetry=False
        )
    )

def get_db_session():
    """获取数据库会话（线程安全）"""
    engine = create_engine(
        f"mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DB}"
    )
    return scoped_session(sessionmaker(bind=engine))
